#include <sxVega_MassSpringSystem.h>
#include <massSpringSystemForceModel.h>
#include <sxVega_EulerIntegrator.h>

Sx::Vega::Physics::VegaMassSpringSystem::VegaMassSpringSystem() {
	this->massSpringSystem = nullptr;
	this->forceModel = nullptr;
	this->massMatrix = nullptr;
	this->integrator = nullptr;
	this->stationary = false;
}

Sx::Vega::Physics::VegaMassSpringSystem::VegaMassSpringSystem(const VegaMassSpringSystem& massSpringSystem) {
	this->massSpringSystem = new Vega_ModifiableMassSpringSystem(*massSpringSystem.massSpringSystem);
	this->construct();
}

Sx::Vega::Physics::VegaMassSpringSystem::~VegaMassSpringSystem() {}

bool Sx::Vega::Physics::VegaMassSpringSystem::constructFromQuadMesh(const std::shared_ptr<Graphics::Objects::Mesh>& mesh, double surfaceDensity, double tensileStiffness, double shearStiffness, double bendStiffness, double damping) {
	if ( mesh == nullptr ) {
		std::cout << "[VegaMassSpringSystem:constructFromQuadMesh] Error: Provided mesh is a nullptr." << std::endl;
		return false;
	}

	unsigned int numParticles = this->vertices.size();
	unsigned int numQuads = mesh->getFaces().size();

	double* restPositions = new double[numParticles * 3];
	int* quads = new int[mesh->getFaces().size() * 4];

	//--------------------------------------------------------------------------
	// Convert the vertices in this physical representation to the rest
	// positions used to construct the VegaFEM mass-spring system.
	//--------------------------------------------------------------------------
	unsigned int index = 0;
	for ( unsigned int i = 0; i < numParticles; i++ ) {
		restPositions[index + 0] = this->vertices[i].x();
		restPositions[index + 1] = this->vertices[i].y();
		restPositions[index + 2] = this->vertices[i].z();
		index += 3;
	}

	//--------------------------------------------------------------------------
	// Convert the face indices from the mesh representation to the int*
	// supported by the VegaFEM mass-spring constructor.
	//--------------------------------------------------------------------------
	index = 0;
	for ( unsigned int i = 0; i < numQuads; i++ ) {
		if ( mesh->getFaces()[i].indices.size() != 4 ) {
			std::cerr << "[VegaMassSpringSystem:constructFromQuadMesh] Error: Mesh contains a face that is not a quad." << std::endl;
			std::cerr << "  All mesh faces must have degree: 4" << std::endl;
			return false;
		}

		quads[index + 0] = mesh->getFaces()[i].indices[0];
		quads[index + 1] = mesh->getFaces()[i].indices[1];
		quads[index + 2] = mesh->getFaces()[i].indices[2];
		quads[index + 3] = mesh->getFaces()[i].indices[3];
		index += 4;
	}

	//--------------------------------------------------------------------------
	// Construct the VegaFEM mass-spring system.
	//--------------------------------------------------------------------------
	this->massSpringSystem = new Vega_ModifiableMassSpringSystem(numParticles, restPositions, numQuads, quads, surfaceDensity, tensileStiffness, shearStiffness, bendStiffness, damping, 0);

	delete [] restPositions;
	delete [] quads;
	return true;
}

bool Sx::Vega::Physics::VegaMassSpringSystem::constructFromTetrahedralMesh(const std::shared_ptr<Graphics::Objects::TetrahedralMesh>& mesh, double density, double tensileStiffness, double damping) {
	if ( mesh == nullptr ) {
		std::cout << "[VegaMassSpringSystem:constructFromTetrahedralMesh] Error: Provided mesh is a nullptr." << std::endl;
		return false;
	}

	unsigned int numElements = mesh->getElements().size();

	if ( numElements == 0 ) {
		std::cerr << "[VegaMassSpringSystem::constructFromTetrahedralMesh] Error: Mesh contains no elements." << std::endl;
		return false;
	}

	unsigned int numParticles = this->vertices.size();
	double* restPositions = new double[numParticles * 3];
	MassSpringSystemElementType elementType = TET;

	unsigned int elementNodeCount = mesh->getElements()[0].indices.size();
	int* elements = new int[numElements * elementNodeCount];

	//--------------------------------------------------------------------------
	// Convert the vertices in this physical representation to the rest
	// positions used to construct the VegaFEM mass-spring system.
	//--------------------------------------------------------------------------
	unsigned int index = 0;
	for ( unsigned int i = 0; i < numParticles; i++ ) {
		restPositions[index + 0] = this->vertices[i].x();
		restPositions[index + 1] = this->vertices[i].y();
		restPositions[index + 2] = this->vertices[i].z();
		index += 3;
	}

	//--------------------------------------------------------------------------
	// Copy the tetrahedral elements from the mesh representation to the
	// element array for the Vega mass spring system.
	//--------------------------------------------------------------------------
	index = 0;
	for ( unsigned int i = 0; i < numElements; i++ ) {
		for ( unsigned int x = 0; x < elementNodeCount; x++ ) {
			elements[index] = mesh->getElements()[i].indices[x];
			index++;
		}
	}

	this->massSpringSystem = new Vega_ModifiableMassSpringSystem(numParticles, restPositions, elementType, numElements, elements, density, tensileStiffness, damping, 0);

	delete [] restPositions;
	delete [] elements;
	return true;
}

bool Sx::Vega::Physics::VegaMassSpringSystem::construct() {
	if ( this->massSpringSystem == nullptr ) {
		std::cerr << "[VegaMassSpringSystem:construct] Error: The VegaFEM mass-spring system must be initialized from a resource." << std::endl;
		std::cerr << "  See any 'constructFromX' functions to initialize the VegaFEM mass-spring system." << std::endl;
		return false;
	}

	this->massSpringSystem->GenerateMassMatrix(&this->massMatrix);
	this->forceModel = new MassSpringSystemForceModel(this->massSpringSystem);
	this->integrator = new Vega_EulerIntegrator(this->massSpringSystem->GetNumParticles(), this->massMatrix, this->forceModel);
	this->integrator->construct(this->vertices);

	this->massSpringSystem->SetGravity(true, -9.81f); // TEMP just to slow it down.
	return true;
}

void Sx::Vega::Physics::VegaMassSpringSystem::update(double dt) {
	if ( this->integrator == nullptr ) return;
	if ( this->massSpringSystem == nullptr ) return;
	if ( this->stationary == true ) return;

	unsigned int particleVectorSize = 3 * this->massSpringSystem->GetNumParticles();

	int integrationCode = 0;
	for ( unsigned int i = 0; i < this->integrator->getSubsteps(); i++ ) {
		this->integrator->performTimestep(dt);

		if ( integrationCode != 0 ) {
			std::cout << "[EnvironmentImp:update] Error: The integrator went unstable. Reduce the timestep, or increase the number of substeps per timestep." << std::endl;
			std::cout << "  Reverting the integrator to utilize initial rest positions." << std::endl;
			this->integrator->reset();
			return;
		}
		
		
	}

	//--------------------------------------------------------------------------
	// Update the prev vertex positions and then update the new vertex
	// positions. These are essentially copy operations.
	//--------------------------------------------------------------------------
	this->updatePreviousVertices();
	this->updateVertexPositions();
}

bool Sx::Vega::Physics::VegaMassSpringSystem::addNodeExternalForce(unsigned int index, Eigen::Vector3d& force) {
	if ( this->integrator == nullptr ) return false;
	bool success = this->integrator->addNodeExternalForce(index, force);
	if ( success == false ) return false;
	if ( index >= this->state.externalForces.size() ) return false;
	this->state.externalForces[index] += force;
	return true;
}

bool Sx::Vega::Physics::VegaMassSpringSystem::addNodeVelocity(unsigned int index, Eigen::Vector3d& velocity) {
	if ( this->integrator == nullptr ) return false;
	bool success = this->integrator->addNodeVelocity(index, velocity);
	if ( success == false ) return false;
	if ( index >= this->state.velocities.size() ) return false;
	this->state.velocities[index] += velocity;
	return true;
}

bool Sx::Vega::Physics::VegaMassSpringSystem::addNodePosition(unsigned int index, Eigen::Vector3d& position) {
	if ( this->integrator == nullptr ) return false;
	bool success = this->integrator->addNodePosition(index, position);
	if ( success == false ) return false;
	if ( index >= this->vertices.size() ) return false;
	this->vertices[index] += position;
	return true;
}

bool Sx::Vega::Physics::VegaMassSpringSystem::addNodeMass(unsigned int index, double mass) {
	if ( this->massSpringSystem == nullptr || this->integrator == nullptr ) return false;
	if ( index >= static_cast<unsigned int>(this->massSpringSystem->GetNumParticles()) ) return false;

	this->massSpringSystem->GetMasses()[index] += mass;
	return true;
}

bool Sx::Vega::Physics::VegaMassSpringSystem::setNodeExternalForce(unsigned int index, Eigen::Vector3d& force) {
	if ( this->integrator == nullptr ) return false;
	bool success = false;
	this->integrator->setNodeExternalForce(index, force);

	if ( success == false ) return false;
	if ( index >= this->state.externalForces.size() ) return false;
	this->state.externalForces[index] = this->integrator->getNodeExternalForce(index);
	return true;
}

bool Sx::Vega::Physics::VegaMassSpringSystem::setNodeExternalForce(unsigned int index, double x, double y, double z) {
	if ( this->integrator == nullptr ) return false;
	bool success = false;
	this->integrator->setNodeExternalForce(index, x, y, z);

	if ( success == false ) return false;
	if ( index >= this->state.externalForces.size() ) return false;
	this->state.externalForces[index] = this->integrator->getNodeExternalForce(index);
	return true;
}

bool Sx::Vega::Physics::VegaMassSpringSystem::setNodeInternalForce(unsigned int index, Eigen::Vector3d& force) {
	if ( this->integrator == nullptr ) return false;
	bool success = false;
	this->integrator->setNodeInternalForce(index, force);

	if ( success == false ) return false;
	if ( index >= this->state.internalForces.size() ) return false;
	this->state.internalForces[index] = this->integrator->getNodeInternalForce(index);
	return true;
}

bool Sx::Vega::Physics::VegaMassSpringSystem::setNodeInternalForce(unsigned int index, double x, double y, double z) {
	if ( this->integrator == nullptr ) return false;
	bool success = false;
	this->integrator->setNodeInternalForce(index, x, y, z);

	if ( success == false ) return false;
	if ( index >= this->state.internalForces.size() ) return false;
	this->state.internalForces[index] = this->integrator->getNodeInternalForce(index);
	return true;
}

bool Sx::Vega::Physics::VegaMassSpringSystem::setNodeVelocity(unsigned int index, Eigen::Vector3d& velocity) {
	if ( this->integrator == nullptr ) return false;
	bool success = false;
	success = this->integrator->setNodeVelocity(index, velocity);

	if ( success == false ) return false;
	if ( index >= this->state.velocities.size() ) return false;
	this->state.velocities[index] = this->integrator->getNodeVelocity(index);
	return true;
}

bool Sx::Vega::Physics::VegaMassSpringSystem::setNodeVelocity(unsigned int index, double x, double y, double z) {
	if ( this->integrator == nullptr ) return false;
	bool success = false;
	success = this->integrator->setNodeVelocity(index, x, y, z);

	if ( success == false ) return false;
	if ( index >= this->state.velocities.size() ) return false;
	this->state.velocities[index] = this->integrator->getNodeVelocity(index);
	return true;
}

bool Sx::Vega::Physics::VegaMassSpringSystem::setNodePosition(unsigned int index, Eigen::Vector3d& position) {
	if ( this->integrator == nullptr ) return false;
	bool success = false;
	success = this->integrator->setNodePosition(index, position);

	if ( success == false ) return false;
	if ( index >= this->vertices.size() ) return false;
	this->vertices[index] = this->integrator->getNodePosition(index);
	return true;
}

bool Sx::Vega::Physics::VegaMassSpringSystem::setNodePosition(unsigned int index, double x, double y, double z) {
	if ( this->integrator == nullptr ) return false;
	bool success = false;
	success = this->integrator->setNodePosition(index, x, y, z);

	if ( success == false ) return false;
	if ( index >= this->vertices.size() ) return false;
	this->vertices[index] = this->integrator->getNodePosition(index);
	return true;
}

bool Sx::Vega::Physics::VegaMassSpringSystem::setNodePreviousPosition(unsigned int index, Eigen::Vector3d& position) {
	if ( index >= this->previousVertices.size() ) return false;
	this->previousVertices[index] = position;
	return true;
}

bool Sx::Vega::Physics::VegaMassSpringSystem::setNodeAcceleration(unsigned int index, Eigen::Vector3d& acceleration) {
	if ( this->integrator == nullptr ) return false;
	bool success = false;
	success = this->integrator->setNodeAcceleration(index, acceleration);

	if ( success == false ) return false;
	if ( index >= this->state.accelerations.size() ) return false;
	this->state.accelerations[index] = this->integrator->getNodeAcceleration(index);
	return true;
}

bool Sx::Vega::Physics::VegaMassSpringSystem::setNodeAcceleration(unsigned int index, double x, double y, double z) {
	if ( this->integrator == nullptr ) return false;
	bool success = false;
	success = this->integrator->setNodeAcceleration(index, x, y, z);

	if ( success == false ) return false;
	if ( index >= this->state.accelerations.size() ) return false;
	this->state.accelerations[index] = this->integrator->getNodeAcceleration(index);
	return true;
}

bool Sx::Vega::Physics::VegaMassSpringSystem::setNodeMass(unsigned int index, double mass) {
	if ( this->massSpringSystem == nullptr || this->integrator == nullptr ) return false;
	if ( index >= static_cast<unsigned int>(this->massSpringSystem->GetNumParticles()) ) return false;

	this->massSpringSystem->GetMasses()[index] = mass;
	return true;
}

void Sx::Vega::Physics::VegaMassSpringSystem::zeroExternalForces() {
	if ( this->integrator == nullptr ) return;
	this->integrator->zeroExternalForces();
	for ( unsigned int i = 0; i < this->state.externalForces.size(); i++ )
		this->state.externalForces[i] = Eigen::Vector3d::Zero();
}

void Sx::Vega::Physics::VegaMassSpringSystem::setSubsteps(unsigned int substeps) {
	if ( this->integrator == nullptr ) return;
	this->integrator->setSubsteps(substeps);
}

bool Sx::Vega::Physics::VegaMassSpringSystem::reset() {
	if ( this->integrator == nullptr ) return false;
	this->integrator->reset();
	for ( unsigned int i = 0; i < this->state.velocities.size(); i++ )
		this->state.velocities[i] = Eigen::Vector3d::Zero();
	for ( unsigned int i = 0; i < this->state.accelerations.size(); i++ )
		this->state.accelerations[i] = Eigen::Vector3d::Zero();
	for ( unsigned int i = 0; i < this->state.externalForces.size(); i++ )
		this->state.externalForces[i] = Eigen::Vector3d::Zero();
	return true;
}

Eigen::Vector3d Sx::Vega::Physics::VegaMassSpringSystem::getNodeExternalForce(unsigned int index) const {
	if ( this->integrator == nullptr ) return Eigen::Vector3d::Zero();
	return this->integrator->getNodeExternalForce(index);
}

Eigen::Vector3d Sx::Vega::Physics::VegaMassSpringSystem::getNodeAcceleraction(unsigned int index) const {
	if ( this->integrator == nullptr ) return Eigen::Vector3d::Zero();
	return this->integrator->getNodeAcceleration(index);
}

Eigen::Vector3d Sx::Vega::Physics::VegaMassSpringSystem::getNodeVelocity(unsigned int index) const {
	if ( this->integrator == nullptr ) return Eigen::Vector3d::Zero();
	return this->integrator->getNodeVelocity(index);
}

Eigen::Vector3d Sx::Vega::Physics::VegaMassSpringSystem::getNodePosition(unsigned int index) const {
	if ( this->integrator == nullptr ) return Eigen::Vector3d::Zero();
	return this->integrator->getNodePosition(index);
}

double Sx::Vega::Physics::VegaMassSpringSystem::getNodeMass(unsigned int index) const {
	if ( this->massSpringSystem == nullptr || this->integrator == nullptr ) return 0;
	if ( index >= static_cast<unsigned int>(this->massSpringSystem->GetNumParticles()) ) return 0;

	return this->massSpringSystem->GetMasses()[index];
}

bool Sx::Vega::Physics::VegaMassSpringSystem::updateVertexPositions() {
	if ( this->massSpringSystem == nullptr ) return false;
	if ( this->integrator == nullptr ) return false;
	
	//--------------------------------------------------------------------------
	// Update the position of each Node.
	//--------------------------------------------------------------------------
	for ( int i = 0; i < this->massSpringSystem->GetNumParticles(); i++ ) {
		this->vertices[i] = this->integrator->getNodePosition(i);
	}

	//--------------------------------------------------------------------------
	// Update the Velocities of all Nodes
	//--------------------------------------------------------------------------
	for ( int i = 0; i < this->massSpringSystem->GetNumParticles(); i++ ) {
		this->state.velocities[i] = this->integrator->getNodeVelocity(i);
	}

	//--------------------------------------------------------------------------
	// Update the accelerations of all Nodes
	//--------------------------------------------------------------------------
	for ( int i = 0; i < this->massSpringSystem->GetNumParticles(); i++ ) {
		this->state.accelerations[i] = this->integrator->getNodeAcceleration(i);
	}

	//--------------------------------------------------------------------------
	// Update the external force for each Node.
	//--------------------------------------------------------------------------
	for ( int i = 0; i < this->massSpringSystem->GetNumParticles(); i++ ) {
		this->state.externalForces[i] = this->integrator->getNodeExternalForce(i);
	}

	for ( int i = 0; i < this->massSpringSystem->GetNumParticles(); i++ ) {
		this->state.internalForces[i] = this->integrator->getNodeInternalForce(i);
	}

	return true;
}
